{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Research_work(TPU support)",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Linar23/Research_work/blob/master/Research_work(TPU_support).ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "BcPg_3DemUbz",
        "colab": {}
      },
      "source": [
        "import os\n",
        "assert os.environ['COLAB_TPU_ADDR'], 'Make sure to select TPU from Edit > Notebook settings > Hardware accelerator'"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "outputId": "152c5aaa-7389-4131-fbb0-19019b344478",
        "id": "Ms52C8UBmUA4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 650
        }
      },
      "source": [
        "DIST_BUCKET=\"gs://tpu-pytorch/wheels\"\n",
        "TORCH_WHEEL=\"torch-1.15-cp36-cp36m-linux_x86_64.whl\"\n",
        "TORCH_XLA_WHEEL=\"torch_xla-1.15-cp36-cp36m-linux_x86_64.whl\"\n",
        "TORCHVISION_WHEEL=\"torchvision-0.3.0-cp36-cp36m-linux_x86_64.whl\"\n",
        "\n",
        "# Install Colab TPU compat PyTorch/TPU wheels and dependencies\n",
        "!pip uninstall -y torch torchvision\n",
        "!gsutil cp \"$DIST_BUCKET/$TORCH_WHEEL\" .\n",
        "!gsutil cp \"$DIST_BUCKET/$TORCH_XLA_WHEEL\" .\n",
        "!gsutil cp \"$DIST_BUCKET/$TORCHVISION_WHEEL\" .\n",
        "!pip install \"$TORCH_WHEEL\"\n",
        "!pip install \"$TORCH_XLA_WHEEL\"\n",
        "!pip install \"$TORCHVISION_WHEEL\"\n",
        "!sudo apt-get install libomp5"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Uninstalling torch-1.3.0a0+111da77:\n",
            "  Successfully uninstalled torch-1.3.0a0+111da77\n",
            "Uninstalling torchvision-0.3.0:\n",
            "  Successfully uninstalled torchvision-0.3.0\n",
            "Copying gs://tpu-pytorch/wheels/torch-1.15-cp36-cp36m-linux_x86_64.whl...\n",
            "-\n",
            "Operation completed over 1 objects/73.3 MiB.                                     \n",
            "Copying gs://tpu-pytorch/wheels/torch_xla-1.15-cp36-cp36m-linux_x86_64.whl...\n",
            "- [1 files][116.6 MiB/116.6 MiB]                                                \n",
            "Operation completed over 1 objects/116.6 MiB.                                    \n",
            "Copying gs://tpu-pytorch/wheels/torchvision-0.3.0-cp36-cp36m-linux_x86_64.whl...\n",
            "/ [1 files][  3.1 MiB/  3.1 MiB]                                                \n",
            "Operation completed over 1 objects/3.1 MiB.                                      \n",
            "Processing ./torch-1.15-cp36-cp36m-linux_x86_64.whl\n",
            "\u001b[31mERROR: fastai 1.0.59 requires torchvision, which is not installed.\u001b[0m\n",
            "Installing collected packages: torch\n",
            "Successfully installed torch-1.3.0a0+111da77\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "torch"
                ]
              }
            }
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: torch-xla==1.15 from file:///content/torch_xla-1.15-cp36-cp36m-linux_x86_64.whl in /usr/local/lib/python3.6/dist-packages (1.15+xrt)\n",
            "Processing ./torchvision-0.3.0-cp36-cp36m-linux_x86_64.whl\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torchvision==0.3.0) (1.17.3)\n",
            "Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.6/dist-packages (from torchvision==0.3.0) (4.3.0)\n",
            "Requirement already satisfied: torch>=1.1.0 in /usr/local/lib/python3.6/dist-packages (from torchvision==0.3.0) (1.3.0a0+111da77)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from torchvision==0.3.0) (1.12.0)\n",
            "Requirement already satisfied: olefile in /usr/local/lib/python3.6/dist-packages (from pillow>=4.1.1->torchvision==0.3.0) (0.46)\n",
            "Installing collected packages: torchvision\n",
            "Successfully installed torchvision-0.3.0\n",
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "libomp5 is already the newest version (5.0.1-1).\n",
            "0 upgraded, 0 newly installed, 0 to remove and 28 not upgraded.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ms7MwO1vs2IR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "import time\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "import torch_xla\n",
        "import torch_xla.distributed.data_parallel as dp\n",
        "import torch_xla.utils.utils as xu\n",
        "import torch_xla.core.xla_model as xm\n",
        "import torch_xla.debug.metrics as met\n",
        "\n",
        "\n",
        "batch_size = 2 #@param {type:\"integer\"}\n",
        "num_workers = 4 #@param {type:\"integer\"}\n",
        "learning_rate = 0.1 #@param {type:\"number\"}\n",
        "momentum = 0.9 #@param {type:\"number\"}\n",
        "num_epochs = 20 #@param {type:\"integer\"}\n",
        "num_cores = 8 #@param [8, 1] {type:\"raw\"}\n",
        "log_steps = 20 #@param {type:\"integer\"}\n",
        "metrics_debug = False #@param {type:\"boolean\"}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IzoeKBt8tEl6",
        "colab_type": "text"
      },
      "source": [
        "## Run"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l1Xtub1dEvsC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!mkdir /content/Graph"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "wr7axV9KcLNT",
        "outputId": "4055b0a3-499f-4f49-c66c-cf4103a2ca92",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 102
        }
      },
      "source": [
        "!pip install tensorboardcolab\n",
        "from tensorboardcolab import TensorBoardColab\n",
        "tb = TensorBoardColab(graph_path='/content/Graph')"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: tensorboardcolab in /usr/local/lib/python3.6/dist-packages (0.0.22)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "Using TensorFlow backend.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Wait for 8 seconds...\n",
            "TensorBoard link:\n",
            "http://60363ae6.ngrok.io\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2DY73jpzL9lv",
        "colab_type": "code",
        "outputId": "6b03c4ee-c90d-435c-cba9-afdd3349814f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 122
        }
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Go to this URL in a browser: https://accounts.google.com/o/oauth2/auth?client_id=947318989803-6bn6qk8qdgf4n4g3pfee6491hc0brc4i.apps.googleusercontent.com&redirect_uri=urn%3Aietf%3Awg%3Aoauth%3A2.0%3Aoob&scope=email%20https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fdocs.test%20https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fdrive%20https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fdrive.photos.readonly%20https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fpeopleapi.readonly&response_type=code\n",
            "\n",
            "Enter your authorization code:\n",
            "··········\n",
            "Mounted at /content/gdrive\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OtJYjM03BFRK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!rm -r -f /content/Graph.zip"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NrnrVl-d-6Ro",
        "colab_type": "code",
        "outputId": "592b2a4e-7574-4ecd-cda6-7f1bf2b674d9",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 647
        }
      },
      "source": [
        "import torch\n",
        "import random\n",
        "from collections import Counter\n",
        "from math import sqrt\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.optim import Adam\n",
        "from sklearn.model_selection import train_test_split\n",
        "import math\n",
        "import numpy as np\n",
        "import pickle\n",
        "\n",
        "\n",
        "class BERTDataset:\n",
        "    def __init__(self, corpus_path, vocab, seq_len):\n",
        "        \"\"\"\n",
        "        :param corpus_path:\n",
        "        :param vocab:\n",
        "        :param seq_len:\n",
        "        \"\"\"\n",
        "        self.vocab = vocab\n",
        "        self.seq_len = seq_len\n",
        "\n",
        "        self.corpus_path = corpus_path\n",
        "\n",
        "        with open(corpus_path, \"r\") as f:\n",
        "            self.lines = [line.replace(\"\\n\", \"\").split(\"\\\\t\") for line in f]\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.lines)\n",
        "\n",
        "    def __getitem__(self, item):\n",
        "        t1, t2, is_next_label = self.random_sent(item)\n",
        "        t1_random, t1_label = self.random_word(t1)\n",
        "        t2_random, t2_label = self.random_word(t2)\n",
        "\n",
        "        t1 = [self.vocab.cls_index] + t1_random + [self.vocab.sep_index]\n",
        "        t2 = t2_random + [self.vocab.sep_index]\n",
        "\n",
        "        t1_label = [self.vocab.pad_index] + t1_label + [self.vocab.pad_index]\n",
        "        t2_label = t2_label + [self.vocab.pad_index]\n",
        "\n",
        "        segment_label = ([1 for _ in range(len(t1))] + [2 for _ in range(len(t2))])[:self.seq_len]\n",
        "        bert_input = (t1 + t2)[:self.seq_len]\n",
        "        bert_label = (t1_label + t2_label)[:self.seq_len]\n",
        "\n",
        "        padding = [self.vocab.pad_index for _ in range(self.seq_len - len(bert_input))]\n",
        "\n",
        "        bert_input += padding\n",
        "        bert_label += padding\n",
        "        segment_label += padding\n",
        "\n",
        "        output = {\"bert_input\": bert_input,\n",
        "                  \"bert_label\": bert_label,\n",
        "                  \"segment_label\": segment_label,\n",
        "                  \"is_next\": is_next_label}\n",
        "\n",
        "        return {key: torch.tensor(value) for key, value in output.items()}\n",
        "\n",
        "    def random_word(self, sentence):\n",
        "        tokens = sentence.split()\n",
        "        output_label = []\n",
        "\n",
        "        for i, token in enumerate(tokens):\n",
        "            prob = random.random()\n",
        "            if prob < 0.15:\n",
        "                prob /= 0.15\n",
        "\n",
        "                if prob < 0.8:\n",
        "                    tokens[i] = self.vocab.mask_index\n",
        "                elif prob < 0.9:\n",
        "                    tokens[i] = random.randrange(len(self.vocab))\n",
        "                else:\n",
        "                    tokens[i] = self.vocab.token_to_index.get(token, self.vocab.unk_index)\n",
        "\n",
        "                output_label.append(self.vocab.token_to_index.get(token, self.vocab.unk_index))\n",
        "            else:\n",
        "                tokens[i] = self.vocab.token_to_index.get(token, self.vocab.unk_index)\n",
        "                output_label.append(self.vocab.pad_index)\n",
        "\n",
        "        return tokens, output_label\n",
        "\n",
        "    def random_sent(self, index):\n",
        "        t1, t2 = self.get_corpus_line(index)\n",
        "\n",
        "        if random.random() > 0.5:\n",
        "            return t1, t2, 1\n",
        "        else:\n",
        "            return t1, self.get_random_line(), 0\n",
        "\n",
        "    def get_corpus_line(self, item):\n",
        "        return self.lines[item][0], self.lines[item][1]\n",
        "\n",
        "    def get_random_line(self):\n",
        "        return self.lines[random.randrange(len(self.lines))][1]\n",
        "\n",
        "\n",
        "class Vocab:\n",
        "    def __init__(self, text):\n",
        "        \"\"\"\n",
        "        :param text:\n",
        "        \"\"\"\n",
        "        self.specials = [\"<pad>\", \"<unk>\", \"<sep>\", \"<cls>\", \"<mask>\"]\n",
        "\n",
        "        self.pad_index = 0\n",
        "        self.unk_index = 1\n",
        "        self.sep_index = 2\n",
        "        self.cls_index = 3\n",
        "        self.mask_index = 4\n",
        "\n",
        "        self.index_to_token = list(self.specials)\n",
        "\n",
        "        counter = Counter()\n",
        "        \n",
        "        for line in text:\n",
        "            words = line.replace(\"\\n\", \"\").replace(\"\\\\t\", \"\").split()\n",
        "\n",
        "            for word in words:\n",
        "                counter[word] += 1\n",
        "\n",
        "        words_and_frequencies = sorted(counter.items())\n",
        "\n",
        "        for word, freq in words_and_frequencies:\n",
        "            if(freq > 200):\n",
        "                self.index_to_token.append(word)\n",
        "\n",
        "        self.token_to_index = {token: i for i, token in enumerate(self.index_to_token)}\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.index_to_token)\n",
        "\n",
        "\n",
        "class ScaledDotProductAttention(nn.Module):\n",
        "    def __init__(self, d_k):\n",
        "        \"\"\"\n",
        "        :param d_k: int scaling factor\n",
        "        \"\"\"\n",
        "        super(ScaledDotProductAttention, self).__init__()\n",
        "\n",
        "        self.scaling = 1 / (sqrt(d_k))\n",
        "\n",
        "    def forward(self, q, k, v, mask):\n",
        "        \"\"\"\n",
        "        :param q:  An float tensor with shape of [b_s, seq_len, d_model / n_head]\n",
        "        :param k: An float tensor with shape of [b_s, seq_len, d_model / n_head]\n",
        "        :param v: An float tensor with shape of [b_s, seq_len, d_model / n_head]\n",
        "\n",
        "        :return: An float tensor with shape of [b_s, seq_len, d_model / n_head]\n",
        "        \"\"\"\n",
        "        attention = torch.bmm(q, k.transpose(1, 2)) * self.scaling\n",
        "        attention = attention.masked_fill(mask == 0, -1e9)\n",
        "\n",
        "        attention = F.softmax(attention, dim=2)\n",
        "\n",
        "        output = torch.bmm(attention, v)\n",
        "\n",
        "        return output\n",
        "\n",
        "\n",
        "class SingleHeadAttention(nn.Module):\n",
        "    def __init__(self, d_model, d_k, d_v):\n",
        "        \"\"\"\n",
        "        :param d_model: Int\n",
        "        :param d_k: Int = d_model / n_head\n",
        "        :param d_v: Int = d_model / n_head\n",
        "        \"\"\"\n",
        "        super(SingleHeadAttention, self).__init__()\n",
        "\n",
        "        self.q_linear = nn.Linear(d_model, d_k)\n",
        "        self.k_linear = nn.Linear(d_model, d_k)\n",
        "        self.v_linear = nn.Linear(d_model, d_v)\n",
        "\n",
        "        self.attention = ScaledDotProductAttention(d_k)\n",
        "\n",
        "    def forward(self, q, k, v, mask):\n",
        "        \"\"\"\n",
        "        :param q: An float tensor with shape of [b_s, seq_len, d_model]\n",
        "        :param k: An float tensor with shape of [b_s, seq_len, d_model]\n",
        "        :param v: An float tensor with shape of [b_s, seq_len, d_model]\n",
        "\n",
        "        :return: An float tensor with shape of [b_s, seq_len, d_model / n_heads]\n",
        "        \"\"\"\n",
        "        proj_q = self.q_linear(q)\n",
        "        proj_k = self.k_linear(k)\n",
        "        proj_v = self.v_linear(v)\n",
        "\n",
        "        output = self.attention(proj_q, proj_k, proj_v, mask)\n",
        "\n",
        "        return output\n",
        "\n",
        "\n",
        "class MultiHeadAttention(nn.Module):\n",
        "    def __init__(self, n_head, d_model):\n",
        "        \"\"\"\n",
        "        :param n_head: Int number of heads\n",
        "        :param d_model: Int\n",
        "        \"\"\"\n",
        "        super(MultiHeadAttention, self).__init__()\n",
        "\n",
        "        d_v = int(d_model / n_head)\n",
        "        d_k = int(d_model / n_head)\n",
        "\n",
        "        self.attention = nn.ModuleList([SingleHeadAttention(d_model, d_k, d_v) for _ in range(n_head)])\n",
        "\n",
        "        self.Linear = nn.Linear(n_head * d_v, d_model)\n",
        "\n",
        "    def forward(self, q, k, v, mask):\n",
        "        \"\"\"\n",
        "        :param q: An float tensor with shape of [b_s, seq_len, d_model]\n",
        "        :param k: An float tensor with shape of [b_s, seq_len, d_model]\n",
        "        :param v: An float tensor with shape of [b_s, seq_len, d_model]\n",
        "\n",
        "        :return: An float tensor with shape of [b_s, seq_len, d_model]\n",
        "        \"\"\"\n",
        "        results = []\n",
        "\n",
        "        for i, single_attention in enumerate(self.attention):\n",
        "            attention_out = single_attention(q, k, v, mask)\n",
        "            results.append(attention_out)\n",
        "\n",
        "        concat = torch.cat(results, dim=2)\n",
        "\n",
        "        linear_output = self.Linear(concat)\n",
        "\n",
        "        return linear_output\n",
        "\n",
        "\n",
        "class TokenEmbedding(nn.Embedding):\n",
        "    def __init__(self, vocab_size, emb_size):\n",
        "        super().__init__(vocab_size, emb_size)\n",
        "\n",
        "\n",
        "class SegmentEmbedding(nn.Embedding):\n",
        "    def __init__(self, emb_size):\n",
        "        super().__init__(3, emb_size)\n",
        "\n",
        "\n",
        "class PositionalEmbedding(nn.Module):\n",
        "    def __init__(self, d_model, max_len=512):\n",
        "        super().__init__()\n",
        "\n",
        "        pe = torch.zeros(max_len, d_model).float()\n",
        "        pe.require_grad = False\n",
        "\n",
        "        position = torch.arange(0, max_len).float().unsqueeze(1)\n",
        "        div_term = torch.pow(10000, torch.arange(0, d_model, 2).float() / d_model)\n",
        "\n",
        "        pe[:, 0::2] = torch.sin(position / div_term)\n",
        "        pe[:, 1::2] = torch.cos(position / div_term)\n",
        "\n",
        "        self.pe = pe.unsqueeze(0)\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.pe[:, :x.size(1)]\n",
        "\n",
        "\n",
        "class BERTEmbedding(nn.Module):\n",
        "    def __init__(self, vocab_size, emb_size):\n",
        "        \"\"\"\n",
        "        :param vocab_size: Int size of vocabulary\n",
        "        :param emb_size: Int size of embedding\n",
        "        \"\"\"\n",
        "        super(BERTEmbedding, self).__init__()\n",
        "\n",
        "        self.v_s = vocab_size\n",
        "        self.e_s = emb_size\n",
        "\n",
        "        self.token = TokenEmbedding(self.v_s, self.e_s)\n",
        "        self.segment = SegmentEmbedding(self.e_s)\n",
        "        self.position = PositionalEmbedding(self.e_s)\n",
        "\n",
        "    def forward(self, seq, segment_label):\n",
        "        \"\"\"\n",
        "        :param seq: An long tensor with shape of [b_s, seq_len]\n",
        "\n",
        "        :return: An float tensor with shape of [b_s, seq_len, emb_size]\n",
        "        \"\"\"\n",
        "\n",
        "        return self.token(seq) + self.segment(segment_label) + self.position(seq)\n",
        "\n",
        "\n",
        "class GELU(nn.Module):\n",
        "    def forward(self, x):\n",
        "        return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))\n",
        "\n",
        "\n",
        "class PositionWise(nn.Module):\n",
        "    def __init__(self, size, inner_size):\n",
        "        \"\"\"\n",
        "        :param size: Int input size\n",
        "        :param inner_size: Int inner size of position wise\n",
        "        \"\"\"\n",
        "        super(PositionWise, self).__init__()\n",
        "\n",
        "        self.fc = nn.Sequential(\n",
        "            nn.Linear(size, inner_size),\n",
        "            GELU(),\n",
        "            nn.Linear(inner_size, size)\n",
        "        )\n",
        "\n",
        "        self.layer_norm = nn.LayerNorm(size)\n",
        "\n",
        "    def forward(self, input):\n",
        "        \"\"\"\n",
        "        :param input: An float tensor with shape of [b_s, seq_len, emb_size]\n",
        "\n",
        "        :return: An float tensor with shape of [b_s, seq_len, emb_size]\n",
        "        \"\"\"\n",
        "        residual = input\n",
        "\n",
        "        result = self.fc(input)\n",
        "\n",
        "        return self.layer_norm(result + residual)\n",
        "\n",
        "\n",
        "class Encoder(nn.Module):\n",
        "    def __init__(self, embeddings, d_model, n_heads, n_layers, vocab_s):\n",
        "        \"\"\"\n",
        "        :param embeddings: An float embeddings tensor with shape [b_s, seq_len, d_model]\n",
        "        :param d_model: Int size of input\n",
        "        :param n_heads: Int number of heads\n",
        "        :param vocab_s: Int size of vocabulary\n",
        "        \"\"\"\n",
        "        super(Encoder, self).__init__()\n",
        "\n",
        "        self.embeddings = embeddings\n",
        "        self.vocab_s = vocab_s\n",
        "\n",
        "        self.transformer_blocks = nn.ModuleList([nn.Sequential(MultiHeadAttention(n_heads, d_model), nn.LayerNorm(d_model), PositionWise(d_model, d_model * 4)) for _ in range(n_layers)])\n",
        "\n",
        "\n",
        "    def forward(self, x, segment_label):\n",
        "        \"\"\"\n",
        "        :param input: An long tensor with shape of [b_s, seq_len]\n",
        "\n",
        "        :return: An float tensor with shape of [b_s, seq_len, vocab_size]\n",
        "        \"\"\"\n",
        "\n",
        "        mask = (x > 0).unsqueeze(1).repeat(1, x.size(1), 1)\n",
        "        input = self.embeddings(x, segment_label)\n",
        "\n",
        "        for multi_head_block, layer_norm, position_wise in self.transformer_blocks:\n",
        "            input = layer_norm(input + multi_head_block(q=input, k=input, v=input, mask=mask))\n",
        "            input = position_wise(input)\n",
        "\n",
        "        return input\n",
        "\n",
        "\n",
        "class Model(nn.Module):\n",
        "    def __init__(self, n_heads, n_layers, vocab_size, emb_size):\n",
        "        \"\"\"\n",
        "        :param n_heads: Int number of heads\n",
        "        :param vocab_size: Int size of vocabulary\n",
        "        :param emb_size: Int embedding size\n",
        "        \"\"\"\n",
        "        super(Model, self).__init__()\n",
        "\n",
        "        self.embed = BERTEmbedding(vocab_size, emb_size)\n",
        "\n",
        "        self.d_model = self.embed.e_s\n",
        "        self.v_s = self.embed.v_s\n",
        "\n",
        "        self.encoder = Encoder(self.embed, self.d_model, n_heads, n_layers, self.v_s)\n",
        "        self.next_sentence = NextSentencePrediction(self.d_model)\n",
        "        self.mask_lm = MaskedLanguageModel(self.d_model, self.v_s)\n",
        "\n",
        "    def forward(self, x, segment_label):\n",
        "        \"\"\"\n",
        "        :param x: An float tensor with shape of [b_s, seq_len]\n",
        "        :param segment_label: An float tensor with shape of [b_s, seq_len]\n",
        "        \"\"\"\n",
        "        prediction = self.encoder(x, segment_label)\n",
        "\n",
        "        return self.next_sentence(prediction), self.mask_lm(prediction)\n",
        "\n",
        "\n",
        "class NextSentencePrediction(nn.Module):\n",
        "    def __init__(self, d_model):\n",
        "        \"\"\"\n",
        "        :param d_model: Int\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.linear = nn.Linear(d_model, 2)\n",
        "        self.softmax = nn.LogSoftmax(dim=-1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.softmax(self.linear(x[:, 0]))\n",
        "\n",
        "\n",
        "class MaskedLanguageModel(nn.Module):\n",
        "    def __init__(self, d_model, vocab_size):\n",
        "        \"\"\"\n",
        "        :param d_model: Int\n",
        "        :param vocab_size: Int size of vocabulary\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.linear = nn.Linear(d_model, vocab_size)\n",
        "        self.softmax = nn.Softmax(dim=-1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.softmax(self.linear(x))\n",
        "\n",
        "\n",
        "class AdamW(torch.optim.Optimizer):\n",
        "    \"\"\"Implements AdamW algorithm.\n",
        "    It has been proposed in `Fixing Weight Decay Regularization in Adam`_.\n",
        "    Arguments:\n",
        "        params (iterable): iterable of parameters to optimize or dicts defining\n",
        "            parameter groups\n",
        "        lr (float, optional): learning rate (default: 1e-3)\n",
        "        betas (Tuple[float, float], optional): coefficients used for computing\n",
        "            running averages of gradient and its square (default: (0.9, 0.999))\n",
        "        eps (float, optional): term added to the denominator to improve\n",
        "            numerical stability (default: 1e-8)\n",
        "        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)\n",
        "    .. Fixing Weight Decay Regularization in Adam:\n",
        "    https://arxiv.org/abs/1711.05101\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,\n",
        "                 weight_decay=0):\n",
        "        defaults = dict(lr=lr, betas=betas, eps=eps,\n",
        "                        weight_decay=weight_decay)\n",
        "        super(AdamW, self).__init__(params, defaults)\n",
        "\n",
        "    def step(self, closure=None):\n",
        "        \"\"\"Performs a single optimization step.\n",
        "        Arguments:\n",
        "            closure (callable, optional): A closure that reevaluates the model\n",
        "                and returns the loss.\n",
        "        \"\"\"\n",
        "        loss = None\n",
        "        if closure is not None:\n",
        "            loss = closure()\n",
        "\n",
        "        for group in self.param_groups:\n",
        "            for p in group['params']:\n",
        "                if p.grad is None:\n",
        "                    continue\n",
        "                grad = p.grad.data\n",
        "                if grad.is_sparse:\n",
        "                    raise RuntimeError('AdamW does not support sparse gradients, please consider SparseAdam instead')\n",
        "\n",
        "                state = self.state[p]\n",
        "\n",
        "                # State initialization\n",
        "                if len(state) == 0:\n",
        "                    state['step'] = 0\n",
        "                    # Exponential moving average of gradient values\n",
        "                    state['exp_avg'] = torch.zeros_like(p.data)\n",
        "                    # Exponential moving average of squared gradient values\n",
        "                    state['exp_avg_sq'] = torch.zeros_like(p.data)\n",
        "\n",
        "                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n",
        "                beta1, beta2 = group['betas']\n",
        "\n",
        "                state['step'] += 1\n",
        "\n",
        "                # according to the paper, this penalty should come after the bias correction\n",
        "                # if group['weight_decay'] != 0:\n",
        "                #     grad = grad.add(group['weight_decay'], p.data)\n",
        "\n",
        "                # Decay the first and second moment running average coefficient\n",
        "                exp_avg.mul_(beta1).add_(1 - beta1, grad)\n",
        "                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)\n",
        "\n",
        "                denom = exp_avg_sq.sqrt().add_(group['eps'])\n",
        "\n",
        "                bias_correction1 = 1 - beta1 ** state['step']\n",
        "                bias_correction2 = 1 - beta2 ** state['step']\n",
        "                step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1\n",
        "\n",
        "                # w = w - wd * lr * w\n",
        "                if group['weight_decay'] != 0:\n",
        "                    p.data.add_(-group['weight_decay'] * group['lr'], p.data)\n",
        "\n",
        "                # w = w - lr * w.grad\n",
        "                p.data.addcdiv_(-step_size, exp_avg, denom)\n",
        "\n",
        "                # w = w - wd * lr * w - lr * w.grad\n",
        "                # See http://www.fast.ai/2018/07/02/adam-weight-decay/\n",
        "\n",
        "        return loss\n",
        "\n",
        "\n",
        "class ScheduledOptim():\n",
        "    '''A simple wrapper class for learning rate scheduling'''\n",
        "\n",
        "    def __init__(self, optimizer, d_model, n_warmup_steps):\n",
        "        self._optimizer = optimizer\n",
        "        self.n_warmup_steps = n_warmup_steps\n",
        "        self.n_current_steps = 0\n",
        "        self.init_lr = np.power(d_model, -0.5)\n",
        "\n",
        "    def step_and_update_lr(self):\n",
        "        \"Step with the inner optimizer\"\n",
        "        self._update_learning_rate()\n",
        "        self._optimizer.step()\n",
        "\n",
        "    def zero_grad(self):\n",
        "        \"Zero out the gradients by the inner optimizer\"\n",
        "        self._optimizer.zero_grad()\n",
        "\n",
        "    def _get_lr_scale(self):\n",
        "        return np.min([\n",
        "            np.power(self.n_current_steps, -0.5),\n",
        "            np.power(self.n_warmup_steps, -1.5) * self.n_current_steps])\n",
        "\n",
        "    def _update_learning_rate(self):\n",
        "        ''' Learning rate scheduling per step '''\n",
        "\n",
        "        self.n_current_steps += 1\n",
        "        lr = self.init_lr * self._get_lr_scale()\n",
        "\n",
        "        for param_group in self._optimizer.param_groups:\n",
        "            param_group['lr'] = lr\n",
        "            \n",
        "            \n",
        "dataset_path = \"/content/gdrive/My Drive/data_small.txt\"\n",
        "\n",
        "with open(dataset_path, \"r\") as f:\n",
        "    vocab = Vocab(f)\n",
        "    \n",
        "#with open(\"/content/vocab_small.pickle\", \"wb\") as f:\n",
        "#    pickle.dump(vocab, f)\n",
        "    \n",
        "# with open(\"/content/dataset_small.pickle\", \"wb\") as f:\n",
        "\n",
        "#with open('/content/vocab_small.pickle', 'rb') as f:\n",
        "#    vocab = pickle.load(f)\n",
        "    \n",
        "#with open('/content/dataset_small.pickle', 'rb') as f:\n",
        "#   dataset = pickle.load(f)\n",
        "\n",
        "\n",
        "seq_len = 64\n",
        "emb_size = 128\n",
        "epochs = 10\n",
        "n_layers = 6\n",
        "n_heads = 3\n",
        "\n",
        "dataset = BERTDataset(dataset_path, vocab, seq_len)\n",
        "\n",
        "train, test = train_test_split(dataset, test_size=0.2)\n",
        "\n",
        "train_data_loader = DataLoader(train, batch_size=batch_size, shuffle=True)\n",
        "print(len(train_data_loader))\n",
        "test_data_loader = DataLoader(test, batch_size=batch_size)\n",
        "#model = torch.load(\"/content/model_epoch_7.pth\").to('cuda')\n",
        "model = Model(n_heads, n_layers, len(vocab), emb_size)\n",
        "\n",
        "masked_criterion = nn.CrossEntropyLoss(ignore_index=0)\n",
        "next_criterion = nn.CrossEntropyLoss()\n",
        "optim = AdamW(model.parameters(), lr=1e-5, weight_decay=0.01)\n",
        "optim_schedule = ScheduledOptim(optim, emb_size, n_warmup_steps=10000)\n",
        "\n",
        "train_loss = []\n",
        "test_loss = []\n",
        "train_acc = []\n",
        "test_acc = []\n",
        "\n",
        "#writer = SummaryWriter()\n",
        "devices = (xm.get_xla_supported_devices(max_devices=num_cores) if num_cores != 0 else [])\n",
        "print(\"Devices: {}\".format(devices))\n",
        "\n",
        "model_parallel = dp.DataParallel(model, device_ids=devices)\n",
        "\n",
        "\n",
        "def train_loop_fn(model, loader, device, context):\n",
        "  masked_criterion = nn.CrossEntropyLoss(ignore_index=0)\n",
        "  next_criterion = nn.CrossEntropyLoss()\n",
        "  optimizer = context.getattr_or('optimizer', lambda: AdamW(model.parameters(), lr=1e-5 * max(len(devices), 1), weight_decay=0.01))\n",
        "  tracker = xm.RateTracker()\n",
        "\n",
        "  model.train()\n",
        "\n",
        "  #total_correct = 0\n",
        "  #total_element = 0\n",
        "\n",
        "  for x, data in enumerate(loader):\n",
        "    optimizer.zero_grad()\n",
        "\n",
        "    next_sent_output, mask_lm_output = model.forward(data[\"bert_input\"], data[\"segment_label\"])\n",
        "\n",
        "    next_loss = next_criterion(next_sent_output, data[\"is_next\"])\n",
        "    mask_loss = masked_criterion(mask_lm_output.transpose(1, 2), data[\"bert_label\"])\n",
        "\n",
        "    loss = next_loss + mask_loss\n",
        "\n",
        "    #correct = next_sent_output.argmax(dim=-1).eq(data[\"is_next\"]).sum().item()\n",
        "    #print(correct)\n",
        "        \n",
        "    #total_correct += correct\n",
        "    #total_element += data[\"is_next\"].nelement()\n",
        "\n",
        "    loss.backward()\n",
        "    xm.optimizer_step(optimizer)\n",
        "    tracker.add(batch_size)\n",
        "    print('[{}]({}) Loss={:.5f} Rate={:.2f} GlobalRate={:.2f} Time={}\\n'.format(\n",
        "          device, x, loss.item(), tracker.rate(),\n",
        "          tracker.global_rate(), time.asctime()), flush=True)\n",
        "    \n",
        "  #print(total_correct * 100.0 / total_element)\n",
        "\n",
        "# Start training\n",
        "\n",
        "for epoch in range(1, num_epochs + 1):\n",
        "  model_parallel(train_loop_fn, train_data_loader)"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "3125\n",
            "Devices: ['xla:1', 'xla:2', 'xla:3', 'xla:4', 'xla:5', 'xla:6', 'xla:7', 'xla:8']\n",
            "[xla:3](0) Loss=9.49090 Rate=4.58 GlobalRate=4.58 Time=Sat Nov  2 11:51:16 2019\n",
            "\n",
            "[xla:8](0) Loss=9.51627 Rate=4.59 GlobalRate=4.59 Time=Sat Nov  2 11:51:16 2019\n",
            "\n",
            "[xla:2](0) Loss=9.57114 Rate=4.57 GlobalRate=4.57 Time=Sat Nov  2 11:51:16 2019\n",
            "\n",
            "[xla:4](0) Loss=9.65058 Rate=4.58 GlobalRate=4.58 Time=Sat Nov  2 11:51:16 2019\n",
            "\n",
            "[xla:1](0) Loss=9.59930 Rate=4.56 GlobalRate=4.56 Time=Sat Nov  2 11:51:16 2019\n",
            "[xla:7](0) Loss=9.56777 Rate=4.59 GlobalRate=4.59 Time=Sat Nov  2 11:51:16 2019\n",
            "[xla:5](0) Loss=9.59177 Rate=4.58 GlobalRate=4.58 Time=Sat Nov  2 11:51:16 2019\n",
            "[xla:6](0) Loss=9.54372 Rate=4.58 GlobalRate=4.58 Time=Sat Nov  2 11:51:16 2019\n",
            "\n",
            "\n",
            "\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-3-c0264a387002>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m    607\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    608\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_epochs\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 609\u001b[0;31m   \u001b[0mmodel_parallel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_loop_fn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain_data_loader\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/torch_xla/distributed/data_parallel.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, loop_fn, loader, fixed_batch_size)\u001b[0m\n\u001b[1;32m    119\u001b[0m       \u001b[0mresults\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    120\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mthread\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mthreads\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 121\u001b[0;31m       \u001b[0mthread\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    122\u001b[0m     \u001b[0mpara_loader\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclose\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    123\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mresult\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mx\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mresults\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/lib/python3.6/threading.py\u001b[0m in \u001b[0;36mjoin\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m   1054\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1055\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mtimeout\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1056\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_wait_for_tstate_lock\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1057\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1058\u001b[0m             \u001b[0;31m# the behavior of a negative timeout isn't documented, but\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/lib/python3.6/threading.py\u001b[0m in \u001b[0;36m_wait_for_tstate_lock\u001b[0;34m(self, block, timeout)\u001b[0m\n\u001b[1;32m   1070\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mlock\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m  \u001b[0;31m# already determined that the C code is done\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1071\u001b[0m             \u001b[0;32massert\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_is_stopped\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1072\u001b[0;31m         \u001b[0;32melif\u001b[0m \u001b[0mlock\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0macquire\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mblock\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtimeout\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1073\u001b[0m             \u001b[0mlock\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrelease\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1074\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_stop\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2ZcQWHAz88H6",
        "colab_type": "code",
        "outputId": "948f87ba-2512-4ca6-ff1c-0a4f57808694",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 322
        }
      },
      "source": [
        "!zip -r /content/Graph.zip /content/Graph\n",
        "\n",
        "from google.colab import files\n",
        "files.download(\"/content/Graph.zip\")"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "updating: content/Graph/ (stored 0%)\n",
            "updating: content/Graph/acc_train/ (stored 0%)\n",
            "updating: content/Graph/avg_loss_train/ (stored 0%)\n",
            "updating: content/Graph/avg_mask_loss_test/ (stored 0%)\n",
            "updating: content/Graph/avg_next_loss_test/ (stored 0%)\n",
            "updating: content/Graph/acc_test/ (stored 0%)\n",
            "updating: content/Graph/avg_next_loss_train/ (stored 0%)\n",
            "updating: content/Graph/avg_mask_loss_train/ (stored 0%)\n",
            "updating: content/Graph/avg_loss_test/ (stored 0%)\n",
            "  adding: content/Graph/acc_train/events.out.tfevents.1559358781.bad86040ca6c (deflated 48%)\n",
            "  adding: content/Graph/avg_loss_train/events.out.tfevents.1559358781.bad86040ca6c (deflated 52%)\n",
            "  adding: content/Graph/avg_mask_loss_test/events.out.tfevents.1559358808.bad86040ca6c (deflated 57%)\n",
            "  adding: content/Graph/avg_next_loss_test/events.out.tfevents.1559358808.bad86040ca6c (deflated 54%)\n",
            "  adding: content/Graph/acc_test/events.out.tfevents.1559358808.bad86040ca6c (deflated 51%)\n",
            "  adding: content/Graph/avg_next_loss_train/events.out.tfevents.1559358781.bad86040ca6c (deflated 55%)\n",
            "  adding: content/Graph/avg_mask_loss_train/events.out.tfevents.1559358781.bad86040ca6c (deflated 56%)\n",
            "  adding: content/Graph/avg_loss_test/events.out.tfevents.1559358808.bad86040ca6c (deflated 52%)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RpiOShbp0rb1",
        "colab_type": "code",
        "outputId": "2843054c-0f60-4b3a-edf2-0b74e6a3ff08",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 306
        }
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Sat Sep 14 11:13:49 2019       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 430.40       Driver Version: 418.67       CUDA Version: 10.1     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla K80           Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   73C    P8    35W / 149W |      0MiB / 11441MiB |      0%      Default |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                       GPU Memory |\n",
            "|  GPU       PID   Type   Process name                             Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TIC67efcffNX",
        "colab_type": "code",
        "outputId": "cdb672bf-7720-4c01-a329-a780b31e7b1d",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 632
        }
      },
      "source": [
        "!pip3 install transformers"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting transformers\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/fd/f9/51824e40f0a23a49eab4fcaa45c1c797cbf9761adedd0b558dab7c958b34/transformers-2.1.1-py3-none-any.whl (311kB)\n",
            "\u001b[K     |████████████████████████████████| 317kB 6.5MB/s \n",
            "\u001b[?25hRequirement already satisfied: tqdm in /usr/local/lib/python3.6/dist-packages (from transformers) (4.28.1)\n",
            "Requirement already satisfied: boto3 in /usr/local/lib/python3.6/dist-packages (from transformers) (1.9.250)\n",
            "Collecting sacremoses (from transformers)\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/1f/8e/ed5364a06a9ba720fddd9820155cc57300d28f5f43a6fd7b7e817177e642/sacremoses-0.0.35.tar.gz (859kB)\n",
            "\u001b[K     |████████████████████████████████| 860kB 44.0MB/s \n",
            "\u001b[?25hCollecting regex (from transformers)\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/ff/60/d9782c56ceefa76033a00e1f84cd8c586c75e6e7fea2cd45ee8b46a386c5/regex-2019.08.19-cp36-cp36m-manylinux1_x86_64.whl (643kB)\n",
            "\u001b[K     |████████████████████████████████| 645kB 37.6MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from transformers) (1.16.5)\n",
            "Collecting sentencepiece (from transformers)\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/14/3d/efb655a670b98f62ec32d66954e1109f403db4d937c50d779a75b9763a29/sentencepiece-0.1.83-cp36-cp36m-manylinux1_x86_64.whl (1.0MB)\n",
            "\u001b[K     |████████████████████████████████| 1.0MB 40.6MB/s \n",
            "\u001b[?25hRequirement already satisfied: requests in /usr/local/lib/python3.6/dist-packages (from transformers) (2.21.0)\n",
            "Requirement already satisfied: jmespath<1.0.0,>=0.7.1 in /usr/local/lib/python3.6/dist-packages (from boto3->transformers) (0.9.4)\n",
            "Requirement already satisfied: botocore<1.13.0,>=1.12.250 in /usr/local/lib/python3.6/dist-packages (from boto3->transformers) (1.12.250)\n",
            "Requirement already satisfied: s3transfer<0.3.0,>=0.2.0 in /usr/local/lib/python3.6/dist-packages (from boto3->transformers) (0.2.1)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers) (1.12.0)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers) (7.0)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers) (0.14.0)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (2019.9.11)\n",
            "Requirement already satisfied: idna<2.9,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (2.8)\n",
            "Requirement already satisfied: urllib3<1.25,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (1.24.3)\n",
            "Requirement already satisfied: chardet<3.1.0,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (3.0.4)\n",
            "Requirement already satisfied: docutils<0.16,>=0.10 in /usr/local/lib/python3.6/dist-packages (from botocore<1.13.0,>=1.12.250->boto3->transformers) (0.15.2)\n",
            "Requirement already satisfied: python-dateutil<3.0.0,>=2.1; python_version >= \"2.7\" in /usr/local/lib/python3.6/dist-packages (from botocore<1.13.0,>=1.12.250->boto3->transformers) (2.5.3)\n",
            "Building wheels for collected packages: sacremoses\n",
            "  Building wheel for sacremoses (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for sacremoses: filename=sacremoses-0.0.35-cp36-none-any.whl size=883999 sha256=d23566e6c01906de8500c525009b0e8961bfc61885057a18d0a09b07a43b393b\n",
            "  Stored in directory: /root/.cache/pip/wheels/63/2a/db/63e2909042c634ef551d0d9ac825b2b0b32dede4a6d87ddc94\n",
            "Successfully built sacremoses\n",
            "Installing collected packages: sacremoses, regex, sentencepiece, transformers\n",
            "Successfully installed regex-2019.8.19 sacremoses-0.0.35 sentencepiece-0.1.83 transformers-2.1.1\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rQoKOUEU0tCB",
        "colab_type": "code",
        "outputId": "59781c52-7567-43f2-c72c-42dc1e308fbc",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 136
        }
      },
      "source": [
        "import torch\n",
        "from transformers import *\n",
        "\n",
        "# Transformers has a unified API\n",
        "# for 8 transformer architectures and 30 pretrained weights.\n",
        "#          Model          | Tokenizer          | Pretrained weights shortcut\n",
        "MODELS = [\n",
        "          (DistilBertForMaskedLM, DistilBertTokenizer, 'distilbert-base-uncased')]\n",
        "\n",
        "# To use TensorFlow 2.0 versions of the models, simply prefix the class names with 'TF', e.g. `TFRobertaModel` is the TF 2.0 counterpart of the PyTorch model `RobertaModel`\n",
        "\n",
        "# Let's encode some text in a sequence of hidden-states using each model:\n",
        "for model_class, tokenizer_class, pretrained_weights in MODELS:\n",
        "    # Load pretrained model/tokenizer\n",
        "    tokenizer = tokenizer_class.from_pretrained(pretrained_weights)\n",
        "    model = model_class.from_pretrained(pretrained_weights)\n",
        "\n",
        "    # Encode text\n",
        "    input_ids = torch.tensor([tokenizer.encode(\"Here is some text to encode\", add_special_tokens=True)])  # Add special tokens takes care of adding [CLS], [SEP], <s>... tokens in the right way for each model.\n",
        "    with torch.no_grad():\n",
        "        last_hidden_states = model(input_ids)  # Models outputs are now tuples\n",
        "        print(last_hidden_states[0])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tensor([[[ -6.0730,  -6.0184,  -6.0532,  ...,  -5.3471,  -5.1327,  -3.2758],\n",
            "         [-11.2493, -11.0514, -11.2565,  ..., -10.8661,  -9.0513,  -7.8534],\n",
            "         [-12.1650, -11.9065, -12.0352,  ..., -10.5904,  -8.4455,  -7.9540],\n",
            "         ...,\n",
            "         [-12.6242, -12.2914, -12.5036,  ..., -11.2858, -10.2073, -11.0556],\n",
            "         [ -8.3184,  -8.1826,  -8.2792,  ...,  -7.9680,  -7.4850,  -8.3284],\n",
            "         [-11.3067, -11.3019, -11.3750,  ...,  -9.1290,  -9.4283,  -7.6095]]])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lG7TBH0XyI3Y",
        "colab_type": "code",
        "outputId": "c8754711-fc4f-4eb6-81bd-a77b046f81d7",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "len(last_hidden_states[0][0][0])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "30522"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 29
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FVGBZlQAf82F",
        "colab_type": "code",
        "outputId": "4a69e337-651f-49a4-a15e-10772a89effe",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "input_ids"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[  101,  2182,  2003,  2070,  3793,  2000,  4372, 16044,   102]])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 30
        }
      ]
    }
  ]
}